import os
import PIL.Image as Image
import torch
from torch.utils.data import Dataset
import torchvision.transforms as transform
import torch.nn as nn
import matplotlib.pyplot as plt

class ImageDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        super().__init__()
        self.image_dir = image_dir
        self.transform = transform
        self.all_images = os.listdir(image_dir)

    def __len__(self):
        return len(self.all_images)

    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.all_images[idx])
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        noise_factor = 0.5
        noisy_images = image + noise_factor*torch.randn(*image.shape)
        noisy_images = torch.clamp(noisy_images, 0., 1.)
        #返回（噪声图片，原始图片）
        return noisy_images,image
    
transform = transform.Compose([
    transform.Resize((68,68)),
    transform.ToTensor()])
full_dataset = ImageDataset(image_dir='../dataset',transform=transform)
print(len(full_dataset))

#划分数据集
train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [0.75, 0.25])

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=32, shuffle=True,drop_last=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=32, shuffle=False)

class ConvDenoiser(nn.Module):
    def __init__(self):
        super(ConvDenoiser, self).__init__()
        #编码层
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(16, 8, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(8, 8, kernel_size=3, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(8, 16, kernel_size=2, stride=2),
            nn.ReLU(),
            nn.ConvTranspose2d(16, 32, kernel_size=2, stride=2),
            nn.ReLU(),
            nn.Conv2d(32, 3, kernel_size=3, padding=1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

model = ConvDenoiser()
# print(model)

x=torch.randn((1,3,68,68))
t= model(x)
print(t.shape)  #torch.Size([1, 3, 68, 68])

loss = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    trian_loss = 0
    for noise_img,tensor_img in train_loader:
        output = model(noise_img)
        loss_value = loss(output, tensor_img)
        optimizer.zero_grad()
        loss_value.backward()
        optimizer.step()
        trian_loss += loss_value.item()
    this_loss = trian_loss / len(train_loader)
    print(f"Epoch [{epoch+1}], Loss: {this_loss:.4f}")

